import numpy as np
from typing import *

class RBM:
    """
    受限玻尔兹曼机参考实现
    """
    def __init__(self, N:int, M:int):
        """
        初始化受限波尔兹曼机
        :param N: 可见层神经元个数
        :param M: 隐藏层神经元个数
        """
        self.W = np.zeros((N, M))
        self.a = np.zeros((N,))
        self.b = np.zeros((M))

    def __sigmoid(self, x:np.ndarray)->np.ndarray:
        """
        非线性激活函数
        :param x: 输入的数组
        :return: 输出的数组
        """
        return 1 / (1 + np.exp(- (x)))

    def __sample(self, p:np.ndarray)->np.ndarray:
        """
        按一定概率采样一个0，1数组
        :param p: 概率向量
        :return: 0/1向量
        """
        values = []
        for i in range(len(p)):
            value = np.random.choice([0,1], p=[1 - p[i], p[i]])
            values.append(value)
        return np.array(values)


    def __sample_hidden(self, v:np.ndarray)->np.ndarray:
        """
        根据的可见状态采样一个隐藏状态
        :param v: 可见状态向量
        :return: 隐藏状态向量
        """
        WT = np.transpose(self.W, [1,0])
        b = self.b
        return self.__sample(self.__sigmoid(WT.dot(v) + b))

    def __sample_visible(self, h:np.ndarray)->np.ndarray:
        """
        根据隐藏状态采样一个可见状态
        :param h: 隐藏状态向量
        :return: 可见状态向量
        """
        W = self.W
        a = self.a
        return self.__sample(self.__sigmoid(W.dot(h) + a))


    def train(self, samples:List[np.ndarray], leaning_rate:float=0.01, epoch:int = 3):
        """
        使用CD-1算法不断优化模型的参数
        :param samples: 训练样本,维度为 (batch-size, N)
        :param leaning_rate: 学习率
        :param epoch: 训练的回合数
        :return: 无
        """
        for i in range(epoch):
            for j in range(len(samples)):
                sample = samples[j]
                h = self.__sample_hidden(sample)
                v1 = self.__sample_visible(h)
                h1 = self.__sample_hidden(v1)
                self.W += leaning_rate * (
                        sample.reshape((len(sample), 1)).dot( h.reshape((1, len(h))) )
                        -
                        v1.reshape((len(v1),1)).dot( h1.reshape((1, len(h1))) )
                )
                self.a += leaning_rate * (sample - v1)
                self.b += leaning_rate * (h - h1)

    def generate(self)->np.ndarray:
        """
        随机生成一个符合模型分布的可观测样本
        :return: 样本
        """
        ramdom_h = [np.random.randint(0,2) for i in range(len(self.b))]
        return self.__sample_visible(np.array(ramdom_h))


